#include "kernel_operator.h"

using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

class KernelLessEqual {
 public:
  __aicore__ inline KernelLessEqual() {}
  __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, 
                              uint32_t blockLength,
                              uint32_t tileNum, uint32_t tileLength,
                              uint32_t lasttileLength, uint32_t tilingKey) {
    ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
    this->blockLength = blockLength;
    this->tileNum =
        tileNum ASSERT(tileNum != 0 && "tile num can not be zero!");
    this->tileLength = tileLength / BUFFER_NUM;
    this->lasttileLength = lasttileLength;

    x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1 + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    x2Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x2 + this->blockLength * GetBlockIdx(),
                        this->blockLength);
    yGm.SetGlobalBuffer(
        (__gm__ uint8_t *)y + this->blockLength * GetBlockIdx(),
        this->blockLength);
  

    this->val = (DTYPE_X1)0;

    pipe.InitBuffer(inQueueIN, BUFFER_NUM, this->tileLength * 2 * sizeof(DTYPE_X1));
    pipe.InitBuffer(outQueueOUT, BUFFER_NUM, this->tileLength * sizeof(uint8_t));
    Init_temp_pip(val);
  }

   __aicore__ inline void Init_temp_pip(half flag) {
    pipe.InitBuffer(tempBuf, this->tileLength * sizeof(half));
    pipe.InitBuffer(tempBuf1, this->tileLength * sizeof(float));
    pipe.InitBuffer(tempBuf2, this->tileLength * sizeof(float));
  }
  
   __aicore__ inline void Init_temp_pip(float flag) {
    pipe.InitBuffer(tempBuf, this->tileLength * sizeof(uint8_t));
    pipe.InitBuffer(tempBuf3, this->tileLength * sizeof(half));
  }

   __aicore__ inline void Init_temp_pip(int32_t flag) {
    pipe.InitBuffer(tempBuf1, this->tileLength * sizeof(float));
    pipe.InitBuffer(tempBuf2, this->tileLength * sizeof(float));
    pipe.InitBuffer(tempBuf, this->tileLength * sizeof(uint8_t));
    pipe.InitBuffer(tempBuf3, this->tileLength * sizeof(float));
    pipe.InitBuffer(tempBuf4, this->tileLength * sizeof(half));
  }

   __aicore__ inline void Init_temp_pip(int8_t flag) {
    pipe.InitBuffer(tempBuf1, this->tileLength * sizeof(half));
    pipe.InitBuffer(tempBuf2, this->tileLength * sizeof(half));
    pipe.InitBuffer(tempBuf, this->tileLength * sizeof(uint8_t));
  }

  __aicore__ inline void Process() {
    int32_t loopCount = this->tileNum * BUFFER_NUM;
    for (int32_t i = 0; i < loopCount; i++) {
      CopyIn(i);
      Compute(i, val);
      CopyOut(i);
    }
  }

 private:
  __aicore__ inline void CopyIn(int32_t progress) {
    LocalTensor<DTYPE_X1> inLocal = inQueueIN.AllocTensor<DTYPE_X1>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
          DataCopy(inLocal[0], x1Gm[0], this->tileLength);
          DataCopy(inLocal[this->tileLength], x2Gm[0], this->tileLength);
        } else {
          //将最后一个分块的起始地址向前移动tileLength-lasttileLength
          DataCopy(
              inLocal[0],
              x1Gm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
          DataCopy(
              inLocal[this->tileLength],
              x2Gm[(progress - 1) * this->tileLength + this->lasttileLength],
              this->tileLength);
        }
      } else {
        DataCopy(inLocal[0], x1Gm[progress * this->tileLength],
                 this->tileLength);
        DataCopy(inLocal[this->tileLength], x2Gm[progress * this->tileLength],
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      //开启double
      //buffer时，由于将输入数据分成了相等的2部分，分块大小为不开启double
      //buffer的一半， 所以需要对最后两个分块数据的起始地址做处理

      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        //分块大小变为tileLength的一半
        //倒数第2个分块数据的起始地址向前移动（tileLength-lasttileLength)，最后一个分块的起始地址以此为基础进行移动
        DataCopy(
            inLocal[0],
            x1Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
        DataCopy(
            inLocal[this->tileLength],
            x2Gm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            (this->tileLength));
      }

      else {
        DataCopy(inLocal[0], x1Gm[progress * (this->tileLength)],
                 (this->tileLength));
        DataCopy(inLocal[this->tileLength], x2Gm[progress * this->tileLength],
                 this->tileLength);
      }
    }

    inQueueIN.EnQue(inLocal);
  }

  __aicore__ inline void Compute(int32_t progress, half flag) {
    LocalTensor<half> inLocal = inQueueIN.DeQue<half>();
    LocalTensor<half> x1Local = inLocal;
    LocalTensor<half> x2Local = inLocal[this->tileLength];

    LocalTensor<float> temp1 = tempBuf1.Get<float>();
    LocalTensor<float> temp2 = tempBuf2.Get<float>();

    LocalTensor<uint8_t> outLocal = outQueueOUT.AllocTensor<uint8_t>();
    LocalTensor<uint8_t> temp = tempBuf.Get<uint8_t>();

    Cast(temp1, x1Local, RoundMode::CAST_NONE, this->tileLength);
    Cast(temp2, x2Local, RoundMode::CAST_NONE, this->tileLength);

    Compare(temp, temp1, temp2, CMPMODE::LE, this->tileLength);
    Sub(temp1, temp1, temp1,this->tileLength);
    Adds(temp1, temp1, (float)1,this->tileLength);

    Select(temp2, temp, temp1, static_cast<float>(0), SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
    Cast(x1Local, temp2, RoundMode::CAST_NONE, this->tileLength);
    Cast(outLocal, x1Local, RoundMode::CAST_NONE, this->tileLength);
    //compare
    outQueueOUT.EnQue<uint8_t>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }

  __aicore__ inline void Compute(int32_t progress, float flag) {
    LocalTensor<float> inLocal = inQueueIN.DeQue<float>();
    LocalTensor<float> x1Local = inLocal;
    LocalTensor<float> x2Local = inLocal[this->tileLength];

    LocalTensor<uint8_t> outLocal = outQueueOUT.AllocTensor<uint8_t>();
    LocalTensor<uint8_t> temp = tempBuf.Get<uint8_t>();
    LocalTensor<half> temp1 = tempBuf3.Get<half>();

    Compare(temp, x1Local, x2Local, CMPMODE::LE, this->tileLength);
    Sub(x1Local,x1Local,x1Local,this->tileLength);
    Adds(x1Local,x1Local,(float)1,this->tileLength);

    Select(x1Local, temp, x1Local, static_cast<float>(0), SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
    Cast(temp1, x1Local, RoundMode::CAST_NONE, this->tileLength);
    Cast(outLocal, temp1, RoundMode::CAST_NONE, this->tileLength);
    //compare
    outQueueOUT.EnQue<uint8_t>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }

  __aicore__ inline void Compute(int32_t progress, int32_t flag) {
    LocalTensor<int32_t> inLocal = inQueueIN.DeQue<int32_t>();
    LocalTensor<int32_t> x1Local = inLocal;
    LocalTensor<int32_t> x2Local = inLocal[this->tileLength];
    LocalTensor<float> tempOneLocal = tempBuf1.Get<float>();
    LocalTensor<float> tempTwoLocal = tempBuf2.Get<float>();
    LocalTensor<uint8_t> outLocal = outQueueOUT.AllocTensor<uint8_t>();
    LocalTensor<uint8_t> temp = tempBuf.Get<uint8_t>();
    LocalTensor<float> temp1 = tempBuf3.Get<float>();
    LocalTensor<half> temp2 = tempBuf4.Get<half>();

    Sub(x2Local, x2Local, x1Local, this->tileLength);
    Sub(x1Local, x1Local, x1Local, this->tileLength);
    Cast(tempOneLocal, x1Local, RoundMode::CAST_RINT, this->tileLength);
    Cast(tempTwoLocal, x2Local, RoundMode::CAST_RINT, this->tileLength);

    Compare(temp, tempOneLocal, tempTwoLocal, CMPMODE::LE, this->tileLength);
    Adds(tempOneLocal,tempOneLocal,(float)1,this->tileLength);

    Select(tempOneLocal, temp, tempOneLocal, static_cast<float>(0), SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
    Cast(temp1, tempOneLocal, RoundMode::CAST_CEIL, this->tileLength);
    Cast(temp2, temp1, RoundMode::CAST_CEIL, this->tileLength);
    Cast(outLocal, temp2, RoundMode::CAST_CEIL, this->tileLength);

    outQueueOUT.EnQue<uint8_t>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }

  __aicore__ inline void Compute(int32_t progress, int8_t flag) {
    LocalTensor<int8_t> inLocal = inQueueIN.DeQue<int8_t>();
    LocalTensor<int8_t> x1Local = inLocal;
    LocalTensor<int8_t> x2Local = inLocal[this->tileLength];
    LocalTensor<uint8_t> temp = tempBuf.Get<uint8_t>();

    LocalTensor<half> tempOneLocal = tempBuf1.Get<half>();
    LocalTensor<half> tempTwoLocal = tempBuf2.Get<half>();
    LocalTensor<uint8_t> outLocal = outQueueOUT.AllocTensor<uint8_t>();

    Cast(tempOneLocal, x1Local, RoundMode::CAST_NONE, this->tileLength);
    Cast(tempTwoLocal, x2Local, RoundMode::CAST_NONE, this->tileLength);

    Compare(temp, tempOneLocal, tempTwoLocal, CMPMODE::LE, this->tileLength);
    Sub(tempOneLocal,tempOneLocal,tempOneLocal,this->tileLength);
    Adds(tempOneLocal,tempOneLocal,(half)1,this->tileLength);

    Select(tempOneLocal, temp, tempOneLocal, static_cast<half>(0), SELMODE::VSEL_TENSOR_SCALAR_MODE, this->tileLength);
    Cast(outLocal, tempOneLocal, RoundMode::CAST_NONE, this->tileLength);

    outQueueOUT.EnQue<uint8_t>(outLocal);

    inQueueIN.FreeTensor(inLocal);
  }

  __aicore__ inline void CopyOut(int32_t progress) {
    LocalTensor<uint8_t> outLocal = outQueueOUT.DeQue<uint8_t>();

    if (BUFFER_NUM == 1) {
      if (progress == this->tileNum - 1) {
        if (progress == 0) {
          //如果只有一包，则搬运的起始地址为0，tileLength为实际分块的数据量
          DataCopy(yGm[0], outLocal, this->tileLength);
        } else {
          //将最后一个分块的起始地址向前移动tileLength-lasttileLength
          DataCopy(
              yGm[(progress - 1) * this->tileLength + this->lasttileLength],
              outLocal, this->tileLength);
        }
      } else {
        DataCopy(yGm[progress * this->tileLength], outLocal,
                 this->tileLength);
      }
    }
    if (BUFFER_NUM == 2) {
      //开启double
      //buffer时，由于将输入数据分成了相等的2部分，分块大小为不开启double
      //buffer的一半， 所以需要对最后两个分块数据的起始地址做处理
      if ((progress == (this->tileNum * BUFFER_NUM - 2)) ||
          (progress == (this->tileNum * BUFFER_NUM - 1))) {
        //分块大小变为tileLength的一半
        //倒数第2个分块数据的起始地址向前移动（tileLength-lasttileLength)，最后一个分块的起始地址以此为基础进行移动
        DataCopy(
            yGm[(progress - 2) * (this->tileLength) + this->lasttileLength],
            outLocal, (this->tileLength));
      }

      else {
        DataCopy(yGm[progress * (this->tileLength)], outLocal,
                 (this->tileLength));
      }
    }

    outQueueOUT.FreeTensor(outLocal);
  }

 private:
  TPipe pipe;
  TQue<QuePosition::VECIN, BUFFER_NUM> inQueueIN;
  TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueOUT;
  GlobalTensor<DTYPE_X1> x1Gm;
  GlobalTensor<DTYPE_X1> x2Gm;
  GlobalTensor<uint8_t> yGm;
  TBuf<> tempBuf;
  TBuf<> tempBuf1;
  TBuf<> tempBuf2;
  TBuf<> tempBuf3;
  TBuf<> tempBuf4;

  DTYPE_X1 val;

  uint32_t blockLength;
  uint32_t tileNum;
  uint32_t tileLength;
  uint32_t lasttileLength;
};

extern "C" __global__ __aicore__ void less_equal(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    KernelLessEqual op;
    uint32_t tilingKey = 1;
    if (TILING_KEY_IS(1)) {
        tilingKey = 1;
    } else if (TILING_KEY_IS(2)) {
        tilingKey = 2;
    } else {
        tilingKey = 1;
    }

    op.Init(x1, x2, y, tiling_data.blockLength,
            tiling_data.tileNum, tiling_data.tileLength,
            tiling_data.lasttileLength, tilingKey);
    op.Process();
}
